import network
import socket
import machine
import time
import array
from machine import I2S, Pin, I2C

# ================= CONFIGURATION =================
SSID = "Astro_RDX"       
PASSWORD = "17060840"    
PC_IP = "192.168.43.151" 
PORT = 50005

# Pin Definitions
PIN_WS = 14      
PIN_SCK = 13     
PIN_SD_MIC = 12  
PIN_SD_SPK = 15  
PIN_BTN_TALK = 2  # 对讲按钮

# === 新增：音量按钮 ===
PIN_VOL_UP = 0    # GP0 上键
PIN_VOL_DOWN = 1  # GP1 下键

PIN_SDA = 26     
PIN_SCL = 27     

SAMPLE_RATE = 16000
CHUNK_SIZE = 1024 
# ===============================================

# --- OLED Init ---
oled = None
try:
    import ssd1306
    import font
    i2c = I2C(1, scl=Pin(PIN_SCL), sda=Pin(PIN_SDA), freq=400000)
    oled = ssd1306.SSD1306_I2C(128, 64, i2c)
    HAS_OLED = True
except:
    HAS_OLED = False

# 全局音量变量 (0-10)
current_vol = 5 
last_status_msg = "Ready" # 记住上一次的状态以便刷新

def draw_ui(status_msg=None):
    """ 统一刷新 UI：上方显示状态，下方显示音量条 """
    global last_status_msg
    if status_msg:
        last_status_msg = status_msg
    
    if not HAS_OLED: return

    try:
        oled.fill(0)
        
        # 1. 顶部状态栏
        if 'font' in globals():
            font.draw(oled, last_status_msg, 0, 0)
        else:
            oled.text(last_status_msg, 0, 0)

        # 2. 底部音量条 (固定位置 y=54)
        # 画一个框
        oled.rect(0, 54, 128, 10, 1) 
        # 画进度条
        bar_width = int((current_vol / 10) * 126)
        oled.fill_rect(2, 56, bar_width, 6, 1)
        
        # 显示数字
        vol_text = f"VOL: {current_vol}"
        oled.text(vol_text, 0, 40) 

        oled.show()
    except: pass

# --- Global Objects ---
led = machine.Pin("LED", machine.Pin.OUT)
btn_talk = machine.Pin(PIN_BTN_TALK, machine.Pin.IN, machine.Pin.PULL_UP)

# === 新增：初始化音量按钮 ===
btn_up = machine.Pin(PIN_VOL_UP, machine.Pin.IN, machine.Pin.PULL_UP)
btn_down = machine.Pin(PIN_VOL_DOWN, machine.Pin.IN, machine.Pin.PULL_UP)

audio_i2s = None 

# --- I2S 管理 (保持之前稳定的逻辑) ---
def init_mic_mode():
    global audio_i2s
    if audio_i2s: audio_i2s.deinit()
    audio_i2s = I2S(0, sck=Pin(PIN_SCK), ws=Pin(PIN_WS), sd=Pin(PIN_SD_MIC),
        mode=I2S.RX, bits=32, format=I2S.MONO, rate=SAMPLE_RATE, ibuf=20000)

def init_spk_mode():
    global audio_i2s
    if audio_i2s: audio_i2s.deinit()
    audio_i2s = I2S(1, sck=Pin(PIN_SCK), ws=Pin(PIN_WS), sd=Pin(PIN_SD_SPK),
        mode=I2S.TX, bits=16, format=I2S.MONO, rate=SAMPLE_RATE, ibuf=40000)

# --- WiFi ---
wlan = network.WLAN(network.STA_IF)
wlan.active(True)
wlan.config(pm=0xa11140) 
wlan.connect(SSID, PASSWORD)
draw_ui("Connecting...")
while not wlan.isconnected():
    time.sleep(0.5)
ip_addr = wlan.ifconfig()[0]
print(f"IP: {ip_addr}")
led.on()

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind(('0.0.0.0', PORT)) 
sock.settimeout(0.05) # 极短超时，保证按键响应灵敏
server_addr = (PC_IP, PORT)
mic_buf = bytearray(CHUNK_SIZE)

init_spk_mode()
current_mode = 'SPK'
draw_ui("System Ready")

# 发送初始音量给电脑
try: sock.sendto(f'VOL:{current_vol}'.encode(), server_addr)
except: pass

# --- 主循环 ---
last_btn_check = 0

try:
    while True:
        # ==========================
        # 1. 音量检测 (任何时候都可用)
        # ==========================
        now = time.ticks_ms()
        if time.ticks_diff(now, last_btn_check) > 150: # 150ms 防抖
            vol_changed = False
            
            # GP0 (UP)
            if btn_up.value() == 0:
                if current_vol < 10:
                    current_vol += 1
                    vol_changed = True
            
            # GP1 (DOWN)
            elif btn_down.value() == 0:
                if current_vol > 0:
                    current_vol -= 1
                    vol_changed = True
            
            if vol_changed:
                last_btn_check = now
                draw_ui() # 刷新屏幕
                # 发送指令给电脑 (格式: VOL:5)
                try: sock.sendto(f'VOL:{current_vol}'.encode(), server_addr)
                except: pass
        
        # ==========================
        # 2. 录音模式 (按住 GP2)
        # ==========================
        if btn_talk.value() == 0:
            if current_mode != 'MIC':
                init_mic_mode()
                current_mode = 'MIC'
                time.sleep(0.05) 

            draw_ui("Recording...")
            led.off()
            try: sock.sendto(b'START', server_addr)
            except: pass
            
            while btn_talk.value() == 0:
                try:
                    num_read = audio_i2s.readinto(mic_buf)
                    if num_read > 0:
                        sock.sendto(mic_buf[:num_read], server_addr)
                except: pass
            
            try: sock.sendto(b'STOP', server_addr)
            except: pass
            
            draw_ui("Thinking...")
            led.on()
            # 清空缓存
            try:
                while True: sock.recv(1024)
            except: pass

            if current_mode != 'SPK':
                init_spk_mode()
                current_mode = 'SPK'

        # ==========================
        # 3. 播放模式 (空闲接收)
        # ==========================
        else:
            if current_mode != 'SPK':
                init_spk_mode()
                current_mode = 'SPK'

            try:
                data, addr = sock.recvfrom(2048)
                if data and len(data) > 0:
                    if data == b'END_OF_STREAM':
                        draw_ui("Ready")
                        # 播静音
                        try: audio_i2s.write(bytearray(1024))
                        except: pass
                    else:
                        # 播放音频
                        if last_status_msg != "Speaking...":
                            draw_ui("Speaking...")
                        audio_i2s.write(data)
            except OSError:
                pass

except KeyboardInterrupt:
    if audio_i2s: audio_i2s.deinit()
    sock.close()
    print("Stopped.")